{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bd1ff156",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d1defc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from scipy.io import loadmat\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from sklearn.metrics import confusion_matrix, classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "920e3c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 数据加载与组织\n",
    "def load_cwru_data(data_path, sample_length=1024, overlap=0.5):\n",
    "    \"\"\"\n",
    "    加载CWRU数据集并进行切片\n",
    "    \n",
    "    参数:\n",
    "    data_path: 数据集路径\n",
    "    sample_length: 每个样本的长度\n",
    "    overlap: 样本重叠率\n",
    "    \n",
    "    返回:\n",
    "    X: 特征数据 [样本数, 样本长度]\n",
    "    y: 标签 [样本数]\n",
    "    \"\"\"\n",
    "    X = []\n",
    "    y = []\n",
    "    \n",
    "    # 定义故障类型和对应的标签\n",
    "    fault_types = {\n",
    "        'normal': 0,\n",
    "        'inner_race': 1,\n",
    "        'outer_race': 2,\n",
    "        'ball': 3\n",
    "    }\n",
    "    \n",
    "    # 遍历不同故障类型的文件夹\n",
    "    for fault_type, label in fault_types.items():\n",
    "        fault_dir = os.path.join(data_path, fault_type)\n",
    "        if not os.path.exists(fault_dir):\n",
    "            continue\n",
    "            \n",
    "        # 遍历该故障类型下的所有文件\n",
    "        for file_name in os.listdir(fault_dir):\n",
    "            if file_name.endswith('.mat'):\n",
    "                file_path = os.path.join(fault_dir, file_name)\n",
    "                mat_data = loadmat(file_path)\n",
    "                \n",
    "                # 提取振动信号\n",
    "                vibration_data = None\n",
    "                for key in mat_data.keys():\n",
    "                    if 'DE_time' in key:  # 驱动端数据\n",
    "                        vibration_data = mat_data[key].flatten()\n",
    "                        break\n",
    "                \n",
    "                if vibration_data is not None:\n",
    "                    # 数据切片\n",
    "                    stride = int(sample_length * (1 - overlap))\n",
    "                    for i in range(0, len(vibration_data) - sample_length + 1, stride):\n",
    "                        X.append(vibration_data[i:i+sample_length])\n",
    "                        y.append(label)\n",
    "    \n",
    "    return np.array(X), np.array(y)\n",
    "\n",
    "# 2. 创建PyTorch数据集\n",
    "class CWRUDataset(Dataset):\n",
    "    def __init__(self, X, y, transform=None):\n",
    "        self.X = torch.tensor(X, dtype=torch.float32).unsqueeze(1)  # 添加通道维度 [B, 1, L]\n",
    "        self.y = torch.tensor(y, dtype=torch.long)\n",
    "        self.transform = transform\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        sample = self.X[idx], self.y[idx]\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "        return sample\n",
    "\n",
    "# 3. 数据预处理管道\n",
    "def preprocess_data(data_path, sample_length=1024, overlap=0.5, test_size=0.2, batch_size=32):\n",
    "    \"\"\"\n",
    "    完整的数据预处理流程\n",
    "    \n",
    "    返回:\n",
    "    train_loader, val_loader, test_loader: 训练、验证和测试数据加载器\n",
    "    class_names: 类别名称列表\n",
    "    \"\"\"\n",
    "    # 加载数据\n",
    "    X, y = load_cwru_data(data_path, sample_length, overlap)\n",
    "    \n",
    "    # 数据标准化\n",
    "    scaler = StandardScaler()\n",
    "    X = scaler.fit_transform(X.reshape(-1, X.shape[1])).reshape(X.shape)\n",
    "    \n",
    "    # 划分训练集、验证集和测试集\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=42, stratify=y)\n",
    "    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42, stratify=y_train)\n",
    "    \n",
    "    # 创建数据集和数据加载器\n",
    "    train_dataset = CWRUDataset(X_train, y_train)\n",
    "    val_dataset = CWRUDataset(X_val, y_val)\n",
    "    test_dataset = CWRUDataset(X_test, y_test)\n",
    "    \n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=batch_size)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "    \n",
    "    class_names = ['正常', '内圈故障', '外圈故障', '滚动体故障']\n",
    "    \n",
    "    return train_loader, val_loader, test_loader, class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82ad6155",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. CBAM注意力机制\n",
    "class ChannelAttention(nn.Module):\n",
    "    def __init__(self, in_channels, reduction_ratio=16):\n",
    "        super(ChannelAttention, self).__init__()\n",
    "        self.avg_pool = nn.AdaptiveAvgPool1d(1)\n",
    "        self.max_pool = nn.AdaptiveMaxPool1d(1)\n",
    "        \n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Conv1d(in_channels, in_channels // reduction_ratio, 1, bias=False),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(in_channels // reduction_ratio, in_channels, 1, bias=False)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        avg_out = self.fc(self.avg_pool(x))\n",
    "        max_out = self.fc(self.max_pool(x))\n",
    "        out = avg_out + max_out\n",
    "        return torch.sigmoid(out)\n",
    "\n",
    "class SpatialAttention(nn.Module):\n",
    "    def __init__(self, kernel_size=7):\n",
    "        super(SpatialAttention, self).__init__()\n",
    "        self.conv = nn.Conv1d(2, 1, kernel_size, padding=kernel_size//2, bias=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        avg_out = torch.mean(x, dim=1, keepdim=True)\n",
    "        max_out, _ = torch.max(x, dim=1, keepdim=True)\n",
    "        out = torch.cat([avg_out, max_out], dim=1)\n",
    "        out = self.conv(out)\n",
    "        return torch.sigmoid(out)\n",
    "\n",
    "# 2. Informer模型的ProbSparse自注意力机制\n",
    "def prob_QK(Q, K, sample_k, n_top):\n",
    "    \"\"\"计算稀疏点的QK值\"\"\"\n",
    "    # Q [B, H, L, D]\n",
    "    B, H, L, E = K.shape\n",
    "    _, _, S, _ = Q.shape\n",
    "    \n",
    "    # 计算K的概率分布\n",
    "    K_expand = K.unsqueeze(-3).expand(B, H, S, L, E)\n",
    "    index_sample = torch.randint(L, (S, sample_k))\n",
    "    K_sample = K_expand[:, :, torch.arange(S).unsqueeze(1), index_sample, :]\n",
    "    \n",
    "    # 计算Q和K_sample的相似度\n",
    "    Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()\n",
    "    \n",
    "    # 找到最相关的n_top个K\n",
    "    M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L)\n",
    "    M_top = M.topk(n_top, sorted=False)[1]\n",
    "    \n",
    "    # 计算这些位置的QK\n",
    "    Q_reduce = Q[torch.arange(B)[:, None, None],\n",
    "                torch.arange(H)[None, :, None],\n",
    "                M_top, :]  # [B, H, n_top, D]\n",
    "    QK = torch.matmul(Q_reduce, K.transpose(-2, -1))  # [B, H, n_top, L]\n",
    "    \n",
    "    return QK, M_top\n",
    "\n",
    "# 3. Informer-CBAM模型\n",
    "class InformerCBAM(nn.Module):\n",
    "    def __init__(self, input_len=1024, d_model=512, n_heads=8, e_layers=3, d_layers=2, \n",
    "                 dropout=0.1, c_out=4, reduction_ratio=16):\n",
    "        super(InformerCBAM, self).__init__()\n",
    "        \n",
    "        # 特征提取层\n",
    "        self.feature_extraction = nn.Sequential(\n",
    "            nn.Conv1d(1, d_model//4, kernel_size=7, stride=2, padding=3),\n",
    "            nn.BatchNorm1d(d_model//4),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=3, stride=2, padding=1),\n",
    "            \n",
    "            nn.Conv1d(d_model//4, d_model//2, kernel_size=5, stride=2, padding=2),\n",
    "            nn.BatchNorm1d(d_model//2),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=3, stride=2, padding=1),\n",
    "            \n",
    "            nn.Conv1d(d_model//2, d_model, kernel_size=3, stride=2, padding=1),\n",
    "            nn.BatchNorm1d(d_model),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        # 计算特征序列长度\n",
    "        seq_len = input_len\n",
    "        for _ in range(3):\n",
    "            seq_len = (seq_len + 1) // 2\n",
    "        \n",
    "        # CBAM注意力机制\n",
    "        self.channel_att = ChannelAttention(d_model, reduction_ratio)\n",
    "        self.spatial_att = SpatialAttention()\n",
    "        \n",
    "        # Informer编码器\n",
    "        self.encoder = nn.ModuleList([\n",
    "            nn.TransformerEncoderLayer(\n",
    "                d_model=d_model,\n",
    "                nhead=n_heads,\n",
    "                dim_feedforward=d_model*4,\n",
    "                dropout=dropout,\n",
    "                activation='gelu',\n",
    "                batch_first=True\n",
    "            )\n",
    "            for _ in range(e_layers)\n",
    "        ])\n",
    "        \n",
    "        # 解码器\n",
    "        self.decoder = nn.ModuleList([\n",
    "            nn.TransformerDecoderLayer(\n",
    "                d_model=d_model,\n",
    "                nhead=n_heads,\n",
    "                dim_feedforward=d_model*4,\n",
    "                dropout=dropout,\n",
    "                activation='gelu',\n",
    "                batch_first=True\n",
    "            )\n",
    "            for _ in range(d_layers)\n",
    "        ])\n",
    "        \n",
    "        # 分类器\n",
    "        self.global_pool = nn.AdaptiveAvgPool1d(1)\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(d_model, d_model//2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(d_model//2, c_out)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x: [B, 1, L]\n",
    "        x = self.feature_extraction(x)  # [B, d_model, seq_len]\n",
    "        \n",
    "        # 应用CBAM注意力\n",
    "        x = x * self.channel_att(x)\n",
    "        x = x * self.spatial_att(x)\n",
    "        \n",
    "        # 转换为Transformer的输入格式 [B, seq_len, d_model]\n",
    "        x = x.transpose(1, 2)\n",
    "        \n",
    "        # 通过编码器\n",
    "        for encoder_layer in self.encoder:\n",
    "            x = encoder_layer(x)\n",
    "        \n",
    "        # 通过解码器（使用编码器输出作为记忆）\n",
    "        for decoder_layer in self.decoder:\n",
    "            x = decoder_layer(x, x)\n",
    "        \n",
    "        # 全局池化和分类\n",
    "        x = self.global_pool(x.transpose(1, 2)).squeeze(-1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "960a5dfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerLSTM(nn.Module):\n",
    "    def __init__(self, input_len=1024, d_model=128, nhead=4, num_layers=2, hidden_size=128, \n",
    "                 dropout=0.1, num_classes=4):\n",
    "        super(TransformerLSTM, self).__init__()\n",
    "        \n",
    "        # 特征提取层\n",
    "        self.feature_extraction = nn.Sequential(\n",
    "            nn.Conv1d(1, d_model//2, kernel_size=7, stride=2, padding=3),\n",
    "            nn.BatchNorm1d(d_model//2),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=3, stride=2, padding=1),\n",
    "            \n",
    "            nn.Conv1d(d_model//2, d_model, kernel_size=5, stride=2, padding=2),\n",
    "            nn.BatchNorm1d(d_model),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=3, stride=2, padding=1)\n",
    "        )\n",
    "        \n",
    "        # 计算特征序列长度\n",
    "        seq_len = input_len\n",
    "        for _ in range(2):\n",
    "            seq_len = (seq_len + 1) // 2\n",
    "        \n",
    "        # Transformer编码器\n",
    "        self.transformer_encoder = nn.TransformerEncoder(\n",
    "            nn.TransformerEncoderLayer(\n",
    "                d_model=d_model,\n",
    "                nhead=nhead,\n",
    "                dim_feedforward=d_model*4,\n",
    "                dropout=dropout,\n",
    "                activation='gelu',\n",
    "                batch_first=True\n",
    "            ),\n",
    "            num_layers=num_layers\n",
    "        )\n",
    "        \n",
    "        # LSTM层\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size=d_model,\n",
    "            hidden_size=hidden_size,\n",
    "            num_layers=num_layers,\n",
    "            batch_first=True,\n",
    "            bidirectional=True,\n",
    "            dropout=dropout if num_layers > 1 else 0\n",
    "        )\n",
    "        \n",
    "        # 分类器\n",
    "        self.global_pool = nn.AdaptiveAvgPool1d(1)\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(hidden_size * 2, hidden_size),  # 双向LSTM输出维度加倍\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(hidden_size, num_classes)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x: [B, 1, L]\n",
    "        x = self.feature_extraction(x)  # [B, d_model, seq_len]\n",
    "        \n",
    "        # 转换为Transformer的输入格式 [B, seq_len, d_model]\n",
    "        x = x.transpose(1, 2)\n",
    "        \n",
    "        # 通过Transformer编码器\n",
    "        x = self.transformer_encoder(x)\n",
    "        \n",
    "        # 通过LSTM层\n",
    "        x, _ = self.lstm(x)\n",
    "        \n",
    "        # 全局池化和分类\n",
    "        x = self.global_pool(x.transpose(1, 2)).squeeze(-1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "074324f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import SVI, Trace_ELBO\n",
    "from pyro.optim import Adam\n",
    "\n",
    "class BayesianNetwork:\n",
    "    def __init__(self, feature_dim, num_classes):\n",
    "        self.feature_dim = feature_dim\n",
    "        self.num_classes = num_classes\n",
    "        self.softmax = torch.nn.Softmax(dim=1)\n",
    "        \n",
    "    def model(self, x, y=None):\n",
    "        # 定义先验分布\n",
    "        w = pyro.sample(\"w\", dist.Normal(torch.zeros(self.feature_dim, self.num_classes), \n",
    "                                        torch.ones(self.feature_dim, self.num_classes)))\n",
    "        b = pyro.sample(\"b\", dist.Normal(torch.zeros(self.num_classes), \n",
    "                                        torch.ones(self.num_classes)))\n",
    "        \n",
    "        # 计算预测概率\n",
    "        logits = torch.matmul(x, w) + b\n",
    "        probs = self.softmax(logits)\n",
    "        \n",
    "        # 定义似然\n",
    "        with pyro.plates(\"data\", x.shape[0]):\n",
    "            obs = pyro.sample(\"obs\", dist.Categorical(probs=probs), obs=y)\n",
    "        return obs\n",
    "    \n",
    "    def guide(self, x, y=None):\n",
    "        # 定义变分分布参数\n",
    "        w_loc = pyro.param(\"w_loc\", torch.randn(self.feature_dim, self.num_classes))\n",
    "        w_scale = pyro.param(\"w_scale\", torch.randn(self.feature_dim, self.num_classes).exp())\n",
    "        b_loc = pyro.param(\"b_loc\", torch.randn(self.num_classes))\n",
    "        b_scale = pyro.param(\"b_scale\", torch.randn(self.num_classes).exp())\n",
    "        \n",
    "        # 采样变分参数\n",
    "        w = pyro.sample(\"w\", dist.Normal(w_loc, w_scale))\n",
    "        b = pyro.sample(\"b\", dist.Normal(b_loc, b_scale))\n",
    "        return w, b\n",
    "    \n",
    "    def train(self, train_loader, num_epochs=100):\n",
    "        # 优化器设置\n",
    "        optimizer = Adam({\"lr\": 0.001})\n",
    "        svi = SVI(self.model, self.guide, optimizer, loss=Trace_ELBO())\n",
    "        \n",
    "        # 训练循环\n",
    "        for epoch in range(num_epochs):\n",
    "            total_loss = 0.0\n",
    "            for x, y in train_loader:\n",
    "                x = x.view(x.size(0), -1)  # 展平为一维特征向量\n",
    "                loss = svi.step(x, y)\n",
    "                total_loss += loss / len(train_loader.dataset)\n",
    "            \n",
    "            if (epoch + 1) % 10 == 0:\n",
    "                print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss:.4f}\")\n",
    "    \n",
    "    def predict(self, x):\n",
    "        x = x.view(x.size(0), -1)  # 展平为一维特征向量\n",
    "        w, b = self.guide(None)\n",
    "        logits = torch.matmul(x, w) + b\n",
    "        probs = self.softmax(logits)\n",
    "        return torch.argmax(probs, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfdbb461",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练函数\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, device, epochs=50):\n",
    "    best_val_acc = 0.0\n",
    "    best_model = None\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        # 训练阶段\n",
    "        model.train()\n",
    "        train_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        \n",
    "        for inputs, targets in train_loader:\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, targets)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            train_loss += loss.item()\n",
    "            _, predicted = outputs.max(1)\n",
    "            total += targets.size(0)\n",
    "            correct += predicted.eq(targets).sum().item()\n",
    "        \n",
    "        train_acc = 100.0 * correct / total\n",
    "        \n",
    "        # 验证阶段\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for inputs, targets in val_loader:\n",
    "                inputs, targets = inputs.to(device), targets.to(device)\n",
    "                \n",
    "                outputs = model(inputs)\n",
    "                loss = criterion(outputs, targets)\n",
    "                \n",
    "                val_loss += loss.item()\n",
    "                _, predicted = outputs.max(1)\n",
    "                total += targets.size(0)\n",
    "                correct += predicted.eq(targets).sum().item()\n",
    "        \n",
    "        val_acc = 100.0 * correct / total\n",
    "        \n",
    "        # 保存最佳模型\n",
    "        if val_acc > best_val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            best_model = model.state_dict().copy()\n",
    "        \n",
    "        print(f'Epoch {epoch+1}/{epochs} | Train Loss: {train_loss/len(train_loader):.4f} | '\n",
    "              f'Train Acc: {train_acc:.2f}% | Val Loss: {val_loss/len(val_loader):.4f} | Val Acc: {val_acc:.2f}%')\n",
    "    \n",
    "    print(f'Best validation accuracy: {best_val_acc:.2f}%')\n",
    "    model.load_state_dict(best_model)\n",
    "    return model\n",
    "\n",
    "# 评估函数\n",
    "def evaluate_model(model, test_loader, device, class_names):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_targets = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for inputs, targets in test_loader:\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            \n",
    "            outputs = model(inputs)\n",
    "            _, predicted = outputs.max(1)\n",
    "            \n",
    "            all_preds.extend(predicted.cpu().numpy())\n",
    "            all_targets.extend(targets.cpu().numpy())\n",
    "    \n",
    "    # 计算混淆矩阵和分类报告\n",
    "    cm = confusion_matrix(all_targets, all_preds)\n",
    "    report = classification_report(all_targets, all_preds, target_names=class_names)\n",
    "    \n",
    "    # 计算总体准确率\n",
    "    accuracy = 100.0 * sum(np.array(all_preds) == np.array(all_targets)) / len(all_preds)\n",
    "    \n",
    "    print(f'Test Accuracy: {accuracy:.2f}%')\n",
    "    print('\\nConfusion Matrix:')\n",
    "    print(cm)\n",
    "    print('\\nClassification Report:')\n",
    "    print(report)\n",
    "    \n",
    "    return accuracy, cm, report\n",
    "\n",
    "# 主函数：运行实验\n",
    "def run_experiment(model_type='informer', data_path='./data/CWRU'):\n",
    "    # 设置设备\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "    # 数据预处理\n",
    "    train_loader, val_loader, test_loader, class_names = preprocess_data(data_path)\n",
    "    \n",
    "    # 创建模型\n",
    "    if model_type == 'informer':\n",
    "        model = InformerCBAM().to(device)\n",
    "    elif model_type == 'transformer_lstm':\n",
    "        model = TransformerLSTM().to(device)\n",
    "    else:\n",
    "        raise ValueError(\"Unsupported model type\")\n",
    "    \n",
    "    # 定义损失函数和优化器\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "    \n",
    "    # 训练模型\n",
    "    print(f\"Training {model_type.upper()} model...\")\n",
    "    model = train_model(model, train_loader, val_loader, criterion, optimizer, device)\n",
    "    \n",
    "    # 评估模型\n",
    "    print(f\"Evaluating {model_type.upper()} model...\")\n",
    "    accuracy, cm, report = evaluate_model(model, test_loader, device, class_names)\n",
    "    \n",
    "    return model, accuracy, cm, report\n",
    "\n",
    "# 贝叶斯网络实验\n",
    "def run_bayesian_experiment(data_path='./data/CWRU'):\n",
    "    # 数据预处理\n",
    "    train_loader, val_loader, test_loader, class_names = preprocess_data(data_path)\n",
    "    \n",
    "    # 获取特征维度\n",
    "    x_sample, _ = next(iter(train_loader))\n",
    "    feature_dim = x_sample.view(x_sample.size(0), -1).size(1)\n",
    "    num_classes = len(class_names)\n",
    "    \n",
    "    # 创建并训练贝叶斯网络\n",
    "    bn = BayesianNetwork(feature_dim, num_classes)\n",
    "    print(\"Training Bayesian Network...\")\n",
    "    bn.train(train_loader, num_epochs=100)\n",
    "    \n",
    "    # 评估贝叶斯网络\n",
    "    print(\"Evaluating Bayesian Network...\")\n",
    "    all_preds = []\n",
    "    all_targets = []\n",
    "    \n",
    "    for x, y in test_loader:\n",
    "        preds = bn.predict(x)\n",
    "        all_preds.extend(preds.numpy())\n",
    "        all_targets.extend(y.numpy())\n",
    "    \n",
    "    # 计算混淆矩阵和分类报告\n",
    "    cm = confusion_matrix(all_targets, all_preds)\n",
    "    report = classification_report(all_targets, all_preds, target_names=class_names)\n",
    "    \n",
    "    # 计算总体准确率\n",
    "    accuracy = 100.0 * sum(np.array(all_preds) == np.array(all_targets)) / len(all_preds)\n",
    "    \n",
    "    print(f'Test Accuracy: {accuracy:.2f}%')\n",
    "    print('\\nConfusion Matrix:')\n",
    "    print(cm)\n",
    "    print('\\nClassification Report:')\n",
    "    print(report)\n",
    "    \n",
    "    return bn, accuracy, cm, report"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
